from rdkit.Chem.SaltRemover import SaltRemover
from math import sqrt
import numpy as np
from collections import defaultdict

# match compatible criteria: |d1 - d2| <= eps1 + eps2 + tau
def check_pair_match_flexible(d1, d2, tau, eps1, eps2):
  if abs(d1 - d2) <= eps1 + eps2 + tau:
    return 1
  else:
    return 0

def check_pair_match(ligand_pair_distance, protein_pair_distance):
  flexibility_constant_tau = 0.6
  interaction_distance_epsilon = 4

  if abs(ligand_pair_distance - protein_pair_distance) <= flexibility_constant_tau + 2 * interaction_distance_epsilon:
    # print(abs(ligand_pair_distance - protein_pair_distance))
    return 1
  else:
    return 0
  
def check_pair_match_with_interation_type(ligand_pair_distance, protein_pair_distance, lig_v_pair):
    flexibility_constant_tau = 0.6
    interaction_distance_epsilon = 4
    v1, v2 = lig_v_pair
    if v1.find('A') != -1 or v1.find('D') != -1 or v2.find('A') != -1 or v2.find('D') != -1: # hydrogen bond interation
        interaction_distance_epsilon = 3
        # print(abs(ligand_pair_distance - protein_pair_distance))
    if abs(ligand_pair_distance - protein_pair_distance) <= flexibility_constant_tau + 2 * interaction_distance_epsilon:
        return 1
    else:
        return 0
    
def check_edge(v1, v2, ligand_distance_dict, protein_distance_dict):
    """check if there is a edgu between BIG vertiex v1 and v2

    Args:
        v1 (tuple): (vl, vp)
        v2 (tuple: (vl, vp) 
    """
    ligand_v_pair = (v1[0], v2[0])
    protein_v_pair = (v1[1], v2[1])

    l_d, p_d = 0, 0
    if ligand_v_pair in ligand_distance_dict:
        l_d = ligand_distance_dict[ligand_v_pair]
    elif ligand_v_pair[::-1] in ligand_distance_dict:
        l_d = ligand_distance_dict[ligand_v_pair[::-1]]
        
    if protein_v_pair in protein_distance_dict:
        p_d = protein_distance_dict[protein_v_pair]
    elif protein_v_pair[::-1] in protein_distance_dict:
        p_d = protein_distance_dict[protein_v_pair[::-1]]

    return  abs(l_d-p_d), check_pair_match_with_interation_type(l_d, p_d, ligand_v_pair)
    

def euclidean_distance(a, b):
  return round(sqrt(sum((a - b)**2 for a, b in zip(a, b))), 1)


def make_dis_matrix(feats):
  init_matrix = np.zeros((len(feats), len(feats)))
  pharmacophore_list = []

  for i in range(len(feats)):
    p1_pos = feats[i].GetPos()
    p1_type = feats[i].GetFamily()
    # print(p1_type)
    pharmacophore_list.append(p1_type)
    for j in range(len(feats)):
      p2_pos = feats[j].GetPos()
      dis = euclidean_distance(p1_pos, p2_pos)
      init_matrix[i][j] = dis
  processed_weight_matrix = init_matrix
  return processed_weight_matrix, pharmacophore_list

def calculate_LDG_vertex_pair_distance(pharmacophore_Abbrev_list, processed_matrix):
# 生成命名顶点对的距离字典
  make_complete_weight_matrix = dict()
  length = processed_matrix.shape[0]
  for i in range(length):
    pair_part1 = pharmacophore_Abbrev_list[i]
    for j in range(i+1, length):
      pair_part2 = pharmacophore_Abbrev_list[j]
    #   print(processed_matrix[i][j])
      make_complete_weight_matrix[(pair_part1, pair_part2)] = processed_matrix[i][j]
  return make_complete_weight_matrix

def find_letter(name:str):
  import re
  pattern = re.compile(r'[a-zA-Z]')   # 查找字母
  result = pattern.findall(name)
  return "".join(result)
# find_letter("Ar1")



def check_pair_vaild(pair):

    allow_pair = {
        "D": "a",
        "d": "A",
        "a": "D",
        "A": "d",
        "n": "P",
        "N": "p",
        "p": "N",
        "P": "n",
        "H": "h",
        "H": "d",
        "h": "H",
        "Ar": "ar",
        "ar": "Ar"
    }

    p1 = pair[0]
    p2 = pair[1]
    pl1 = find_letter(p1)
    pl2 = find_letter(p2)
    if pl2 == allow_pair[pl1]:
        return pair
    else:
        return 0


def make_Abbrev_list(pharmacophore_list: list, moltype: str = "ligand") -> list:
    # 给 LDG 图的顶点命名
    Donor = 0
    Acceptor = 0
    NegIonizable = 0
    PosIonizable = 0
    Aromatic = 0
    LumpedHydrophobe = 0

    pharmacophore_Abbrev_list = []

    if moltype == "ligand":
        for i in pharmacophore_list:
            if i == "Donor":
                name = "D" + str(Donor + 1)
                pharmacophore_Abbrev_list.append(name)
                Donor = Donor + 1
            elif i == 'Acceptor':
                name = "A" + str(Acceptor + 1)
                pharmacophore_Abbrev_list.append(name)
                Acceptor = Acceptor + 1
            elif i == 'NegIonizable':
                name = "N" + str(NegIonizable + 1)
                pharmacophore_Abbrev_list.append(name)
                NegIonizable = NegIonizable + 1
            elif i == 'PosIonizable':
                name = "P" + str(PosIonizable + 1)
                pharmacophore_Abbrev_list.append(name)
                PosIonizable = PosIonizable + 1
            elif i == 'Aromatic':
                name = "Ar" + str(Aromatic + 1)
                pharmacophore_Abbrev_list.append(name)
                Aromatic = Aromatic + 1
            # elif i == 'LumpedHydrophobe':
            elif i == 'Hydrophobe':
                name = "H" + str(LumpedHydrophobe + 1)
                pharmacophore_Abbrev_list.append(name)
                LumpedHydrophobe = LumpedHydrophobe + 1
            else:
                raise Exception("the pharmacophore name is not existed! ")
        return pharmacophore_Abbrev_list

    elif moltype == "receptor":
        for i in pharmacophore_list:
            if i == "Donor":
                name = "d" + str(Donor + 1)
                pharmacophore_Abbrev_list.append(name)
                Donor = Donor + 1
            elif i == 'Acceptor':
                name = "a" + str(Acceptor + 1)
                pharmacophore_Abbrev_list.append(name)
                Acceptor = Acceptor + 1
            elif i == 'NegIonizable':
                name = "n" + str(NegIonizable + 1)
                pharmacophore_Abbrev_list.append(name)
                NegIonizable = NegIonizable + 1
            elif i == 'PosIonizable':
                name = "p" + str(PosIonizable + 1)
                pharmacophore_Abbrev_list.append(name)
                PosIonizable = PosIonizable + 1
            elif i == 'Aromatic':
                name = "ar" + str(Aromatic + 1)
                pharmacophore_Abbrev_list.append(name)
                Aromatic = Aromatic + 1
            # elif i == 'LumpedHydrophobe':
            elif i == 'Hydrophobe':
                name = "h" + str(LumpedHydrophobe + 1)
                pharmacophore_Abbrev_list.append(name)
                LumpedHydrophobe = LumpedHydrophobe + 1
            else:
                raise Exception("the pharmacophore name is not existed! ")
        return pharmacophore_Abbrev_list